# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 20:06 
@Author  : ZhangShenao 
@File    : 1.使用ChatPromptTemplate.py 
@Desc    : 使用ChatPromptTemplate,构造提示词模板
"""
import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate

dotenv.load_dotenv()

# 创建SystemPrompt
system_template = '你是一位专业的顾问，专注于为{product}领域的公司起名字。'
system_prompt = SystemMessagePromptTemplate.from_template(system_template)

# 创建HumanPrompt
human_template = '我们公司的主打产品是{product_detail}，请帮我起2个合适的名字。'
human_prompt = HumanMessagePromptTemplate.from_template(human_template)

# 创建完整Prompt
prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

llm = ChatZhipuAI(model='glm-4-air')
parser = StrOutputParser()
chain = prompt | llm | parser

output = chain.invoke({
    'product': '鲜花装饰',
    'product_detail': '创新的鲜花设计',
})

print(output)
